// Main driver program for simulating condensed-phase burning in a 2D PMMA
// slab.
//
// M. Bruns
// March 2015

#include <dolfin.h>
#include <mshr.h>
#include "HeatTransfer.h"
#include "Velocity.h"
#include "MassLoss.h"
#include "LocalMeshCoarsening.h"

#include <iostream>
#include <fstream>

using namespace dolfin;

// ------------------------------------------------------------------------
// Supporting Functions and Classes
// ------------------------------------------------------------------------

// Boundary classes

class SymmetryBoundary : public SubDomain
{
    bool inside(const Array<double>& x, bool on_boundary) const
    {
        return x[0] < DOLFIN_EPS;
    }
};

class SolidBoundary : public SubDomain
{
    bool inside(const Array<double>& x, bool on_boundary) const
    {
        return x[1] < DOLFIN_EPS;
    }
};

class GasBoundary : public SubDomain
{

    public:
        
        GasBoundary( double a, double b )
        {
            width = a;
            height = b;
        }

    private:

        double width, height;
        
        bool inside(const Array<double>& x, bool on_boundary) const
        {
            return x[0] > width/2.0 - DOLFIN_EPS or x[1] > height - DOLFIN_EPS;
        }

};

// User defined nonlinear problem
class EnergyEquation : public NonlinearProblem
{
    public:
    
        // Constructor
        EnergyEquation(const Mesh& mesh, 
                       const MeshFunction<std::size_t>& boundary,
                       const Constant& dt,
                       const Constant& T_inf, const Constant& h,
                       const Constant& q_ext, const Constant& k,
                       const Constant& rho, const Constant& c_p,
                       const Constant& dh, const Constant& A,
                       const Constant& E_a, const Constant& eps)
        {
            init<HeatTransfer::FunctionSpace, HeatTransfer::JacobianForm,
                 HeatTransfer::ResidualForm>(mesh, boundary, dt, T_inf, 
                                             h, q_ext, k, rho, c_p, dh, A, 
                                             E_a, eps);
        }
        
        // User defined residual vector
        void F(GenericVector& b, const GenericVector& x)
        {
            
            //std::cout << "Sizes: " << std::endl;
            //std::cout << x.size() << std::endl;
            //std::cout << _T_0->vector()->size() << std::endl;
            //std::cout << _T->vector()->size() << std::endl;
            //
            //std::cout << "Min: " << std::endl;
            //std::cout << x.min() << std::endl;
            //std::cout << _T_0->vector()->min() << std::endl;
            //std::cout << _T->vector()->min() << std::endl;
            //
            //std::cout << "Max: " << std::endl;
            //std::cout << x.max() << std::endl;
            //std::cout << _T_0->vector()->max() << std::endl;
            //std::cout << _T->vector()->max() << std::endl;
           
            //std::cout << "RESIDUAL FORM" << std::endl;
            //std::cout << "Rank: " << (*L).rank() << std::endl;

            // Assemble RHS (Neumann boundary conditions)
            Assembler assembler;
            assembler.assemble(b, *L);
        }
        
        // User defined Jacobian
        void J(GenericMatrix& A, const GenericVector& x)
        {
            // Assemble system
            Assembler assembler;
            assembler.assemble(A, *a);
        }

        // Return solution function
        Function& T()
        { return *_T; }
        
        // Return previous solution function
        Function& T_0()
        { return *_T_0; }

        // Reset forms for coarsened mesh
        void reset_forms(const Mesh& mesh,
                  const MeshFunction<std::size_t>& boundary,
                  const Constant& dt,
                  const Constant& T_inf, const Constant& h,
                  const Constant& q_ext, const Constant& k,
                  const Constant& rho, const Constant& c_p,
                  const Constant& dh, const Constant& A,
                  const Constant& E_a, const Constant& eps)
        {
                
            // Create function space and functions
            std::shared_ptr<HeatTransfer::FunctionSpace> 
                            V(new HeatTransfer::FunctionSpace(mesh));
            _T.reset(new Function(V));
            _T_0.reset(new Function(V));
           
            // Create forms and attach functions
            HeatTransfer::JacobianForm* _a = new HeatTransfer::JacobianForm(V, V);
            HeatTransfer::ResidualForm* _L = new HeatTransfer::ResidualForm(V);
            _a->T = *_T;
            _a->dt = dt; 
            _a->k = k; 
            _a->rho = rho;
            _a->c_p = c_p;
            _a->dh = dh; 
            _a->A = A; 
            _a->E_a = E_a; 
            _a->h = h;
            _a->eps = eps;
            
            _L->T = *_T; _L->T_0 = *_T_0;
            _L->T = *_T;
            _L->dt = dt; 
            _L->T_inf = T_inf; 
            _L->k = k; 
            _L->rho = rho;
            _L->c_p = c_p;
            _L->dh = dh; 
            _L->A = A; 
            _L->E_a = E_a; 
            _L->h = h;
            _L->eps = eps;
            _L->q_ext = q_ext;
           
            // Attach marked boundaries to bilinear and linear forms
            _a->ds = boundary;
            _L->ds = boundary;

            // Wrap pointers in a smart pointer
            a.reset(_a);
            L.reset(_L);

        }

        boost::scoped_ptr<Function> _T;

    private:
     
        template<class X, class Y, class Z>
        void init(const Mesh& mesh,
                  const MeshFunction<std::size_t>& boundary,
                  const Constant& dt,
                  const Constant& T_inf, const Constant& h,
                  const Constant& q_ext, const Constant& k,
                  const Constant& rho, const Constant& c_p,
                  const Constant& dh, const Constant& A,
                  const Constant& E_a, const Constant& eps)
        {

            // Create function space and functions
            std::shared_ptr<X> V(new X(mesh));
            _T.reset(new Function(V));
            _T_0.reset(new Function(V));

            // Create forms and attach functions
            Y* _a = new Y(V, V);
            Z* _L = new Z(V);
            _a->T = *_T;
            _a->dt = dt; 
            _a->k = k; 
            _a->rho = rho;
            _a->c_p = c_p;
            _a->dh = dh; 
            _a->A = A; 
            _a->E_a = E_a; 
            _a->h = h;
            _a->eps = eps;
            
            _L->T = *_T; _L->T_0 = *_T_0;
            _L->T = *_T;
            _L->dt = dt; 
            _L->T_inf = T_inf; 
            _L->k = k; 
            _L->rho = rho;
            _L->c_p = c_p;
            _L->dh = dh; 
            _L->A = A; 
            _L->E_a = E_a; 
            _L->h = h;
            _L->eps = eps;
            _L->q_ext = q_ext;
           
            // Attach marked boundaries to bilinear and linear forms
            _a->ds = boundary;
            _L->ds = boundary;
 
            // Wrap pointers in a smart pointer
            a.reset(_a);
            L.reset(_L);

            // Set temperature to ambient as initial condition
            *_T = T_inf;

        }
     
        // Function space, forms and functions
        boost::scoped_ptr<Form> a;
        boost::scoped_ptr<Form> L;
        //boost::scoped_ptr<Function> _T;
        boost::scoped_ptr<Function> _T_0;
    
};

Function velocity( const Mesh& mesh, const Function& T, 
                    const Constant& g, const Constant& mu,
                    const Constant& rho, const Constant& A,
                    const Constant& E_a )
{
    
    
    // Construct function space
    Velocity::FunctionSpace W(mesh);
    
    // Create forms and mark subspaces
    Velocity::BilinearForm a(W, W);
    Velocity::LinearForm L(W);
    SubSpace V(W, 0);
    SubSpace V_x(V, 0);
    SubSpace V_y(V, 1);
    SubSpace Q(W, 1);

    Function q_int(W);
    Function T_int = q_int[1];
    T_int.interpolate( T );
    
    // Apply coefficients
    a.mu = mu;
    
    L.T = T_int;
    L.g = g;
    L.rho = rho;
    L.A = A;
    L.E_a = E_a;

    // Apply boundary conditions
    SolidBoundary bot_bound;
    SymmetryBoundary sym_bound;
    
    Constant zero_velocity(0.0);
    DirichletBC bc1(V_y, zero_velocity, bot_bound);
    DirichletBC bc2(V_x, zero_velocity, sym_bound);
    std::vector<const DirichletBC*> bcs;
    bcs.push_back(&bc1); bcs.push_back(&bc2);

    // Compute solution
    Function w(W);
    solve(a == L, w, bcs);

    // Extract velocity
    Function u = w[0];
    
    return u;
    
}

// ------------------------------------------------------------------------
// Small Cell Removal Function
// ------------------------------------------------------------------------

void remesh( Mesh& mesh )
{

    BoundaryMesh boundary( mesh, "exterior" );
    double distance = 0.0;
    int idx_origin;

    for (VertexIterator v1(boundary); !v1.end(); ++v1)
    {
        
        std::cout << "Vertex " << v1->index() << " has neighbors:" << std::endl;
        for (VertexIterator v2(v1); !v2.end(); ++v2)
        {
            std::cout << "   " << v2->index() << std::endl;
        }
        std::cout << "_______________________________" << std::endl;        
//        //if ((v->point()).norm() < DOLFIN_EPS)
//        //    idx_origin = v->index();
    }


}


void remove_cells( Mesh& mesh, const MeshFunction<std::size_t>& boundary_parts, 
                   double vol_min )
{

    // Local variables 
    Mesh coarse_mesh(mesh);
    const std::size_t num_cells = mesh.size(mesh.topology().dim());
    const std::size_t num_vertices = mesh.size(0);
    MeshFunction<bool> cell_to_remove(mesh);
    MeshFunction<bool> vertex_to_remove(mesh);
    cell_to_remove.init(mesh.topology().dim());
    vertex_to_remove.init(0);
    std::size_t num_cells_to_remove = 0;
    std::size_t num_vertices_to_remove = 0;
    const CellType& cell_type = mesh.type();
    MeshEditor editor;
    std::size_t vertex;
    std::size_t cv_idx;
    std::size_t num_neighbor_facets;
    std::size_t current_cell = 0;
    std::vector<std::size_t> cell_vertices(cell_type.num_entities(0));
    std::vector<int> old2new_vertex(mesh.num_vertices());
    
    // Compute cells to remove (on exposed boundary and smaller than vol_min)
    vertex_to_remove.set_all( false );

    for (CellIterator c(mesh); !c.end(); ++c)
    {
        
        cell_to_remove[c->index()] = false;
        
        for (FacetIterator f(*c); !f.end(); ++f)
        {
            
            if (boundary_parts.values()[f->index()] == 1 and
                c->volume() < vol_min)
                cell_to_remove[c->index()] = true;
        
        }

        // Check for vertices to be removed
        if (cell_to_remove[c->index()])
        {
            num_cells_to_remove++;
            for (VertexIterator v(*c); !v.end(); ++v)
            {
                num_neighbor_facets = 0;
                for (FacetIterator f(*v); !f.end(); ++f)
                    num_neighbor_facets++;
                if (num_neighbor_facets == mesh.topology().dim())
                    vertex_to_remove[v->index()] = true;
            }

        }

    }

    // Count number of vertices to remove
    vertex = 0;
    for (VertexIterator v(mesh); !v.end(); ++v)
    {

        if (vertex_to_remove[v->index()])
        {
            old2new_vertex[v->index()] = -1;
            num_vertices_to_remove++;
        }
        else
        {
            old2new_vertex[v->index()] = vertex;
            vertex++;
        }

    }
    
    // Print status and exit if no cells need to be removed
    if (num_cells_to_remove > 0)
    {
        std::cout << "Removing small cells..." << std::endl;
        std::cout << "Number of cells in old mesh: " << num_cells << "; to remove: "
                  << num_cells_to_remove << std::endl;
        std::cout << "Number of vertices in old mesh: " << num_vertices << "; to remove: "
                  << num_vertices_to_remove << std::endl;
    }
    else
    {        
        std::cout << "No cells to remove." << std::endl;
        return;
    }
 
    // Prepare mesh editor for coarse_mesh
    editor.open(coarse_mesh, cell_type.cell_type(),
                mesh.topology().dim(), mesh.geometry().dim());
    editor.init_vertices_global(num_vertices - num_vertices_to_remove, 
                                num_vertices - num_vertices_to_remove);
    editor.init_cells_global(num_cells - num_cells_to_remove,
                             num_cells - num_cells_to_remove);
  
    
    // Add old vertices to coarse_mesh
    vertex = 0;
    for (VertexIterator v(mesh); !v.end(); ++v)
    {
        
        if (vertex_to_remove[*v] == false)
        {
            editor.add_vertex(vertex, v->point());
            vertex++;
        }
    }

    // Add old cells to coarse_mesh
    for (CellIterator c(mesh); !c.end(); ++c)
    {
        
        if (cell_to_remove[*c] == false)
        {
            cv_idx = 0;
            for (VertexIterator v(*c); !v.end(); ++v)
                cell_vertices[cv_idx++] = old2new_vertex[v->index()];
            editor.add_cell(current_cell++, cell_vertices);
        }

    }

    editor.close();

    // Update mesh
    mesh = coarse_mesh;
    mesh.order();

}

void update_boundaries( const Mesh& mesh, MeshFunction<std::size_t>& boundary_parts )
{
    
    // Update boundary_parts MeshFunction 
    MeshFunction<std::size_t> boundary_parts_new(mesh, mesh.topology().dim() - 1);
    SymmetryBoundary sym_bound;
    SolidBoundary bot_bound;
   
    bot_bound.mark(boundary_parts_new, 2);
    sym_bound.mark(boundary_parts_new, 3);

    for (FacetIterator f(mesh); !f.end(); ++f)
    {
        if ( f->exterior() and
             boundary_parts_new.values()[f->index()] != 2 and
             boundary_parts_new.values()[f->index()] != 3 )
        {
            boundary_parts_new.set_value( f->index(), 1 );
        }
        
    }

    boundary_parts = boundary_parts_new;

}

// ------------------------------------------------------------------------
// Main Function
// ------------------------------------------------------------------------

int main(int argc, char* argv[])
{
    
    init(argc, argv);

    // ------------------------------------------------------------------------
    // Problem Parameters
    // ------------------------------------------------------------------------
    
    // Numerical parameters
    int N = 64;                     // relative number of cells
    double t_f = 60.0;               // final simulation time, s
    
    Constant dt(1.0e-0);            // time step size, s
    
    // Scenario parameters
    double width = 0.080;               // sample width, m
    double height = 8.5e-3;              // sample thickness, m

    Constant T_inf(298.0);          // ambient temperature, K
    Constant h(10.0);               // heat transfer coefficient, W/m-K
    Constant q_ext(0.0, -52e3);     // external heat flux vector, W/m^2
    Constant g(0.0, -9.81);         // gravitational acceleration, m/s^2

    // Material properties, for PMMA
    Constant k(0.2);                // thermal conductivity, W/m-K
    Constant rho(1100.0);           // density, kg/m^3
    Constant c_p(2200.0);           // specific heat capacity, J/kg-K
    Constant dh(870e3);             // specific heat of gasification, J/kg
    Constant A(8.5e12);             // gasification pre-exponential, 1/s
    Constant E_a(1.88e5);           // gasification activation energy, J/mol
    Constant eps(0.85);             // emissivity
    Constant mu(5e4);               // viscosity, N-s/m^2 (at 200 C) 

    // Local variables
    double t = 0.0;                 // current time, s
    double MLR;                     // mass loss rate, kg/m^2-s
    double mass;                    // mass of system, kg
    double r_0;                     // reference cell radius
    
    Mesh mesh;
    //Mesh new_mesh;

    // ------------------------------------------------------------------------
    // Setup Mesh, Boundaries, and Nonlinear Problem
    // ------------------------------------------------------------------------

    // Generate mesh
    mshr::Rectangle mesh_boundary(Point(0.0, 0.0, 0.0), Point(width/2.0, height, 0.0));
    mshr::generate(mesh, mesh_boundary, N);

    //remesh( mesh );

    // Mark boundaries
    MeshFunction<std::size_t> boundary(mesh, mesh.topology().dim() - 1);
    SymmetryBoundary sym_bound;
    SolidBoundary bot_bound;
    GasBoundary gas_bound(width, height);
    
    gas_bound.mark(boundary, 1);
    bot_bound.mark(boundary, 2);
    sym_bound.mark(boundary, 3);

    // Nonlinear heat equation
    EnergyEquation energy_equation(mesh, boundary, dt, T_inf, h, q_ext, 
                                   k, rho, c_p, dh, A, E_a, eps);

    // ------------------------------------------------------------------------
    // Set up Solver
    // ------------------------------------------------------------------------
    
    // Solution functions
    Function& T = energy_equation.T();
    Function& T_0 = energy_equation.T_0();

    // Create nonlinear solver and set parameters
    NewtonSolver newton_solver;
    newton_solver.parameters["linear_solver"] = "lu";
    newton_solver.parameters["convergence_criterion"] = "incremental";
    newton_solver.parameters["maximum_iterations"] = 10;
    newton_solver.parameters["relative_tolerance"] = 1e-6;
    newton_solver.parameters["absolute_tolerance"] = 1e-15;
    newton_solver.parameters["report"]= false;
    
    // Save initial condition to file
    File file("../../results/energy_equation.pvd", "compressed");
    //file << energy_equation.T();

    // Set up function to compute velocities
    Velocity::FunctionSpace W(mesh);
    Function w1(W);
    Function w2(W);
    Function u = w1[0];
    Function disp = w2[0];
   
    // Set up forms for mass loss rate calculations
    MassLoss::Form_mlr MLR_form(mesh);
    MLR_form.rho = rho;
    MLR_form.A = A;
    MLR_form.E_a = E_a;
    MLR_form.T = T;

    MassLoss::Form_mass mass_form(mesh);
    mass_form.rho = rho;

    // Create file for storing mass loss rate data
    std::ofstream mlr_file;
    mlr_file.open("../../results/mlr.out");
    mlr_file << "time (s), MLR (kg/m^2-s), mass (kg)" << std::endl;

    // critical parameter for checking for cell removal
    double v_crit = 8e-2*width*height;
    r_0 = v_crit/height;

    // Solve
    while (t < t_f)
    {
        
        // Update for next time step
        t += dt;
        std::cout << "time = " << t << std::endl;
        *T_0.vector() = *T.vector();
        
        //newton_solver.solve(energy_equation, *((energy_equation.T()).vector()));
        newton_solver.solve(energy_equation, *T.vector());

        // Compute mass loss rate and mass
        MLR = assemble(MLR_form)*2.0/width;
        mass = assemble(mass_form)*2.0*width;
        
        // Save data to file
        mlr_file << t << "," << MLR << "," << mass << std::endl;
        file << std::pair<const Function*, double>(&T, t);
        
        // Move mesh
        u = velocity( mesh, T, g, mu, rho, A, E_a );
        *disp.vector() = *u.vector();
        *disp.vector() *= dt;
        mesh.move(disp);
        mesh.smooth(3);

//        // TESTING STUFF
//        Mesh new_mesh(mesh);
//        remove_cells( new_mesh, boundary, v_crit );
//        HeatTransfer::FunctionSpace V(new_mesh);
//        Function T_new(V);
//        T_new.interpolate(T);
//
//        std::cout << T_new.vector()->size() << std::endl;
//        std::cout << new_mesh.size(0) << std::endl;
//        std::cout << T_new.vector()->min() << std::endl;
//        std::cout << T_new.vector()->max() << std::endl;
//
//        for (VertexIterator v(new_mesh); !v.end(); ++v)
//        {
//            std::vector<double> vertex_vals;
//            T_new.compute_vertex_values(vertex_vals, new_mesh);
//            
//            std::cout << "---------------------------------------" << std::endl;
//            std::cout << "Vertex index: " << v->index() << std::endl;
//            std::cout << "Vertex coordinates: " << v->point() << std::endl;
//            std::cout << "T_new (getitem): " << T_new.vector()->getitem(v->index()) <<
//                                        std::endl;
//            std::cout << "T_new (vert_vals): " << vertex_vals[v->index()] <<
//                                        std::endl;
//        }
        
//        MeshFunction<std::size_t> new_boundary(new_mesh, new_mesh.topology().dim() - 1);
//        update_boundaries( new_mesh, new_boundary );
//
//        EnergyEquation new_energy_equation(new_mesh, new_boundary, dt, T_inf, h, q_ext, 
//                                        k, rho, c_p, dh, A, E_a, eps);
//    
//        Function& new_T = new_energy_equation.T();
//        Function& new_T_0 = new_energy_equation.T_0();
//
//        *new_T.vector() = *T_new.vector();
//        *new_T_0.vector() = *new_T.vector();
//        
//        // Create nonlinear solver and set parameters
//        NewtonSolver new_newton_solver;
//        new_newton_solver.parameters["linear_solver"] = "lu";
//        new_newton_solver.parameters["convergence_criterion"] = "incremental";
//        new_newton_solver.parameters["maximum_iterations"] = 10;
//        new_newton_solver.parameters["relative_tolerance"] = 1e-6;
//        new_newton_solver.parameters["absolute_tolerance"] = 1e-15;
//        new_newton_solver.parameters["report"]= false;
//
//        std::cout << "hai" << std::endl;
//        new_newton_solver.solve(new_energy_equation, *new_T.vector());
//        
//        plot(new_T);
//        interactive();

        

        // --------------------------
        
//        // Remove cells
//        if (mesh.rmin() < r_0)
//        {
//            
//            new_mesh = mesh;
//            
//            remove_cells( new_mesh, boundary, v_crit );
//
//            HeatTransfer::FunctionSpace V(new_mesh);
//            Function T_new(V);
//            
//            
//            T_new.interpolate(energy_equation.T());
//
//            mesh = new_mesh;
//            update_boundaries( mesh, boundary );
//
//            energy_equation.reset_forms(mesh, boundary, dt, T_inf, h, q_ext, 
//                                        k, rho, c_p, dh, A, E_a, eps);
//
//            *energy_equation._T->vector() = *T_new.vector();
//
//
//        }
//
    }

    // Plot solution
    plot(energy_equation.T());
    interactive();

    // Close file
    mlr_file.close();

    return 0;

}
